import os
import io
from tempo import *
from tempo0 import *  # 含有setParameters等函数
import base64
import pywt
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import time


def add_noise_to_trajectory(trajectory, noise_std):  # 加噪函数
    """
    向航迹数据中添加高斯噪声。

    参数:
    trajectory (numpy.ndarray): 形状为 (N, D) 的航迹数据，其中 N 是点的数量，D 是每个点的维度。
    noise_std (float): 噪声的标准差。

    返回:
    numpy.ndarray: 添加了噪声的航迹数据。
    """
    # 生成高斯噪声
    noise = np.random.normal(0, noise_std, size=trajectory.shape)
    # 添加噪声
    noisy_trajectory = trajectory + noise
    return noisy_trajectory


def wavelet_denoise(data,wavelet='haar', level=4, mode='symmetric'):
    """
    使用小波变换对数据进行去噪。
    sym4  haar  db4
    参数:
    data (numpy.ndarray): 待去噪的数据。
    wavelet (str): 使用的小波基函数，默认为 'db4'。
    level (int): 小波分解的层数，默认为 1。
    mode (str): 边界处理模式，默认为 'symmetric'。

    返回:
    numpy.ndarray: 去噪后的数据。
    """

    def sure_shrink(coeffs, noise_std):
        """
        自适应阈值选择：SUREShrink
        :param coeffs: 小波变换得到的系数列表
        :param noise_std: 噪声的标准差估计
        :return: 优化后的系数列表
        """
        def soft_threshold(coeff, threshold):
            return np.sign(coeff) * np.maximum(np.abs(coeff) - threshold, 0)
        
        # 初始化最优阈值和风险
        best_risk = np.inf
        best_threshold = None
        
        # 对每个尺度的系数进行自适应阈值选择
        for detail_level in range(1, len(coeffs)):
            # 初始阈值
            initial_threshold = noise_std * np.sqrt(2 * np.log(len(coeffs[detail_level])))
            
            # 在一定范围内搜索最佳阈值
            thresholds = np.linspace(0, 2 * initial_threshold, 100)
            risks = []
            
            for threshold in thresholds:
                # 应用软阈值处理
                coeffs_shrinked = soft_threshold(coeffs[detail_level], threshold)
                
                # 计算 SURE 风险估计
                risk = np.sum(coeffs_shrinked ** 2) + 2 * threshold * np.sum(np.abs(coeffs_shrinked)) - noise_std ** 2 * len(coeffs[detail_level])
                risks.append(risk)
            
            # 寻找最小风险对应的阈值
            min_risk_idx = np.argmin(risks)
            best_threshold = thresholds[min_risk_idx]
            
            # 更新系数
            coeffs[detail_level] = soft_threshold(coeffs[detail_level], best_threshold)

        return coeffs

    # level = pywt.dwt_max_level(len(data), pywt.Wavelet(wavelet).dec_len)

    coeffs = pywt.wavedec(data, wavelet, level=level,mode=mode)

    # 计算噪声标准差的估计
    mad = lambda x: np.median(np.abs(x - np.median(x))) / 0.6745
    sigma = mad(coeffs[-level])

    noise_std = sigma
    # 对每个尺度的系数进行自适应阈值处理
    coeffs_denoised = sure_shrink(coeffs, noise_std)
    # 进行小波逆变换
    data_denoised = pywt.waverec(coeffs_denoised, wavelet,mode=mode)
    # 确保输出长度与输入长度一致
    if len(data_denoised) != len(data):
        diff = len(data) - len(data_denoised)
        if diff > 0:
            data_denoised = np.pad(data_denoised, (0, diff), mode='edge')
        elif diff < 0:
            data_denoised = data_denoised[:len(data)]    
    return data_denoised


def plot_n_trajectory(groups, labels, colors, show=True, save_path=None):
    """Plot multiple trajectories on a single plot."""
    fig = plt.figure(figsize=(20, 4))
    axs = [fig.add_subplot(1, 3, i + 1, projection=('3d' if i == 0 else None)) for i in range(3)]

    for group, label, color in zip(groups, labels, colors):
        # Convert timestamp to datetime
        times_dt = pd.to_datetime(group['time'], unit='s')
        times_num = mdates.date2num(times_dt)

        # Extract values from the group
        latitudes = group.lat.values
        longitudes = group.lon.values
        geoAltitudes = group.geoaltitude.values
        times = group.time.values

        # Plot the 3D trajectory
        axs[0].scatter3D(longitudes, latitudes, geoAltitudes, s=2, color=color, label=label)

        # Plot the 2D scatter plot
        axs[1].scatter(longitudes, latitudes, s=2, color=color, label=label)

        # Plot the altitude plot
        axs[2].plot(times_num, geoAltitudes, color=color, label=label)

    # Set titles for the columns
    axs[0].set_title('3D trajectory')
    axs[1].set_title('2D scatter plot')
    axs[2].set_title('Altitude plot')

    axs[0].set_xlabel('Longitude')
    axs[0].set_ylabel('Latitude')
    axs[0].set_zlabel('GeoAltitude')

    axs[1].set_xlabel('Longitude')
    axs[1].set_ylabel('Latitude')

    axs[2].set_xlabel('Time')
    axs[2].set_ylabel('GeoAltitude')
    axs[2].set_xticks(axs[2].get_xticks()[::2])  # keep every other tick

    # Set the x-axis as time
    for ax in axs[2:]:
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%H:%M:%S'))  # format the time as HH:MM:SS

    # Add legends to the plots
    for ax in axs:
        ax.legend()

    # Show the plot if required
    if show:
        plt.tight_layout()
        plt.show()

    # Save the plot if save_path is provided
    if save_path is not None:
        fig.savefig(save_path)

    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    image_base64 = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close()
    return image_base64


def main():
    # 记录程序开始的时间
    start_time = time.time()
    # 输入参数提取
    parameters = getParameters()
    # 功能一：导入数据
    folder_path = parameters["folder_path"]                     # string类型，数据导入的文件夹路径，文件夹应含有多条npy文件
    # num=parameters["hangji_number"]                           # int类型，导入航迹的数量，取前hangji_number条
    selected_number = parameters["selected_number"]             # int类型，选择一个文件序号
    
    # 功能二：加噪
    noise_std = parameters["noise_std"]                         # float类型，加噪标准差，高斯白噪声

    # 功能三：去噪
    wavelet_type = parameters["wavelet_type"]                   # string类型，小波变换基函数的类型，做下拉菜单，可选sym4  haar  db4
    level = parameters["level"]                                 # int类型，小波分解层数，一般四层左右效果较好
    mode = parameters["mode"]                                   # string类型，小波变换边界处理函数，做下拉菜单，可选symmetric
                                                                # periodic reflect constant zero smooth antireflect
    denoised_save_path = parameters["denoised_save_path"]       # string类型，去噪航迹保存的文件夹路径

    '''功能一：加载航迹数据'''
   
    file_name = os.listdir(folder_path)
    # 只取前 n 个文件名
    # file_name0 = file_name[:num]
    lines = []  # 二维string类型数组，保存所有导入的航迹全路径名和本文件名，如[r'C:\Users\22377\Desktop\Variflight_CCA1919_20210316.xls3.npy',
                # 'Variflight_CCA1919_20210316.xls3.npy']
    for alphabet in file_name:
        # 加载 .npy 文件                         
        lines.append([folder_path+'\\'+alphabet, alphabet])

    # 读入所有数据
    data = []

    for i in range(len(lines)):
        line = np.load(lines[i][0])
        data.append(line)

        # 仅加载时间、经纬度、高度
    
    trajectory = data[selected_number-1]
    '''trajectory0= np.zeros((len(trajectory),len(trajectory[0])))
    trajectory0[:,0]=trajectory[:,1]
    trajectory0[:,1]=trajectory[:,2]'''
    # 可视化原始数据

    '''# 输出一:原始航迹'''
    plt.figure(figsize=(12, 6))
    plt.plot(trajectory[:, 1], trajectory[:, 2], label='Original Trajectory')
    plt.legend()
    plt.xlabel('Latitude')
    plt.ylabel('Longitude')
    plt.title('Original Track')
    plt.grid(True)
    plt.show()
    plt.close()
    # 返回图像
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    image_base64_original = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close()

    '''功能二：加噪'''
    # 创建一个空列表来存储非时间列的数据
    noisy_data = []
    # 遍历每个轨迹数据集
    # 为每个航迹数据集添加噪声，并与时间列合并
    for i in range(len(data)):
        # 提取非时间列（第 1、2、3 列）
        non_time_data = data[i][:, 1:4]
        
        # 为轨迹数据添加噪声
        noisy_trajectory0 = add_noise_to_trajectory(non_time_data, noise_std)
        
        # 提取时间列
        time_column = data[i][:, 0]
        
        # 合并时间列和加噪后的数据
        combined_data = np.column_stack((time_column, noisy_trajectory0))
        
        # 将合并后的数据添加到 noisy_data 列表中
        noisy_data.append(combined_data)

    noisy_trajectory = noisy_data[selected_number-1]   # 选中的加噪航迹

    # 可视化原始数据和添加噪声后的数据
    '''# 输出二:原始航迹+噪声航迹'''
    plt.figure(figsize=(12, 6))
    plt.plot(trajectory[:, 1], trajectory[:, 2], label='Original Trajectory')
    plt.plot(noisy_trajectory[:, 1], noisy_trajectory[:, 2], label='Noisy Trajectory', linestyle='--')
    plt.legend()
    plt.xlabel('Latitude')
    plt.ylabel('Longitude')
    plt.title('Original、Noised')
    plt.grid(True)
    plt.show()
    plt.close()
    # 返回图像输出二图像编码字符串
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    image_base64_original_noised = base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close()
    '''# 如果文件夹不存在，则创建它
    if not os.path.exists(noised_save_path):
        os.makedirs(noised_save_path)
    # 保存添加噪声后的航迹数据到新的 .npy 文件
    np.save(noised_save_path+"\\"+no_noise_path+'_noiesd.npy', noisy_trajectory)'''

    '''功能三：去噪'''
    
    # 创建一个与原始数据相同形状的数组，用于存放去噪后的数据
    denoised_data = []

    # 对每一列数据进行去噪
    for i in range(len(noisy_data)):
        denoise_data = np.zeros_like(noisy_data[i])
        for col in range(noisy_data[i].shape[1]):
            if col == 0:  # 忽略时间列
                continue
            denoise_data[:, col] = wavelet_denoise(noisy_data[i][:, col], wavelet_type, level, mode)
            denoise_data[:, 0] = noisy_data[i][:, 0]
            denoised_data.append(denoise_data)

    denoised_trajectory = denoised_data[selected_number-1]
    print(denoised_trajectory)
    # 可视化原始数据、加噪数据和去噪后的数据
    '''# 输出三:原始航迹+噪声航迹+去噪航迹'''
    plt.figure(figsize=(12, 6))
    plt.plot(trajectory[:, 1], trajectory[:, 2], label='Original Trajectory')
    plt.plot(noisy_trajectory[:, 1], noisy_trajectory[:, 2], label='noised Trajectory', linestyle='-')
    plt.plot(denoised_trajectory[:, 1], denoised_trajectory[:, 2], label='Denoised Trajectory', linestyle='--')
    plt.title('Original、Noised、Denoised')
    plt.xlabel('Latitude')
    plt.ylabel('Longitude')
    plt.legend()
    plt.grid(True)
    plt.show()
    # 返回图像
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    image_base64_original_noised_denoised=base64.b64encode(buf.getvalue()).decode('utf-8')
    plt.close()

    '''输出四'''
    # 保存去噪文件到指定的文件夹
    # 如果文件夹不存在，则创建它
    if not os.path.exists(denoised_save_path):
        os.makedirs(denoised_save_path)
    # 保存去除噪声后的航迹数据到新的 .npy 文件
    np.save(denoised_save_path+"\\"+lines[selected_number-1][1]+'_denoiesd.npy', denoised_trajectory)

    '''输出五：计算去噪误差'''
    # 示例数据
    original_coords = trajectory
    denoised_coords = denoised_trajectory
    # 计算每个时间点的去噪误差（使用欧几里得距离）
    errors = []
    for original, denoised in zip(original_coords, denoised_coords):
        # 提取经纬度
        original_longitude, original_latitude = original[1], original[2]
        denoised_longitude, denoised_latitude = denoised[1], denoised[2]
        # 计算欧几里得距离
        error = np.sqrt((original_longitude - denoised_longitude)**2 + (original_latitude - denoised_latitude)**2)
        errors.append(error)
    # 转换为 NumPy 数组以便后续计算
    errors = np.array(errors)
    # 计算均方误差（Mean Squared Error, MSE）
    mse = np.mean(errors**2)
    
    # 计算平均误差
    avg_error = np.mean(errors)
    
    # 计算最大误差
    max_error = np.max(errors)

    # 计算中位数误差
    median_error = np.median(errors)

    columns = ['time', 'lat', 'lon', 'geoaltitude']
    traj_df = pd.DataFrame(trajectory, columns=columns)
    noise_traj_df = pd.DataFrame(noisy_trajectory, columns=columns)
    denoised_df = pd.DataFrame(denoised_trajectory, columns=columns)
    image_base64_original = plot_n_trajectory([traj_df], ['Original'], ['red'])
    image_base64_original_noised = plot_n_trajectory([traj_df, noise_traj_df], ['Original', 'Noise'], ['red', 'green'])
    image_base64_original_noised_denoised = plot_n_trajectory([traj_df, noise_traj_df, denoised_df],
                                                              ['Original', 'Noise', 'Denoised'], ['red', 'green', 'blue'])
    
    # 记录程序结束的时间
    end_time = time.time()
    # 计算程序运行时间
    execution_time = end_time - start_time
    # 输出
    '''输出一:原始航迹'''
    setShowData({"image_base64": image_base64_original})                         # 生成的原始航迹图， base64 Png格式的图片
    '''输出二:原始航迹+噪声航迹'''
    setShowData({"image_base64": image_base64_original_noised})                  # 生成的原始航迹和加噪航迹图， base64 Png格式的图片
    '''输出三:原始航迹+噪声航迹+去噪航迹'''
    setShowData({"image_base64": image_base64_original_noised_denoised})         # 生成的原始航迹、加噪航迹、去噪航迹图， base64 Png格式的图片
    '''输出四'''
    # 保存去噪文件到指定的文件夹，但是应该不用设置接口
    '''输出五:估计误差'''
    setShowData({"eval_result": {"运行时间": execution_time, "均方误差": mse, "平均误差": avg_error, "最大误差": max_error,
                                 "中位数误差": median_error}})


if __name__ == "__main__":
    # 数据输入部分
    parameters = {
        # 功能一：导入数据
        "folder_path": r'D:\pythonProject\粒子滤波\去噪新\经纬高归一数据',     # string类型，数据导入的文件夹路径，文件夹应含有多条npy文件
        "selected_number": 1,                                       # int类型，选择一个文件序号(1-n)
        
        # 功能二：加噪
        "noise_std": 0.01,                                          # float类型，加噪标准差，高斯白噪声

        # 功能三：去噪
        "wavelet_type": 'haar',                                      # string类型，小波变换基函数的类型，做下拉菜单，可选sym4  haar  db4
        "level": 4,                                                  # int类型，小波分解层数，一般四层左右效果较好
        "mode": 'symmetric',                                         # string类型，小波变换边界处理函数，做下拉菜单，
                                                                     # 可选symmetric periodic reflect constant zero
                                                                     # smooth antireflect
        "denoised_save_path": r'D:\pythonProject\粒子滤波\去噪新\去噪航迹'       # string类型，去噪航迹保存的文件夹路径

    }
    setParameters(parameters)

    main()
